import sys
sys.path.append("..")
from SentModel.Sent2Vec import *
from Dataloader.twitter_loader import *
from SentModel.sentence_trainer import *
from torch.utils.data import DataLoader

tr = LMReader("../data/twitter_LM_tr.txt")
dev = LMReader("../data/twitter_LM_dev.txt")
te = LMReader("../data/twitter_LM_te.txt")

batchsize_GPU = 16
GPU_Count = torch.cuda.device_count()
train_batchsize = batchsize_GPU*GPU_Count
tr_loader = DataLoader(tr, batch_size=train_batchsize, shuffle=True)
dev_loader = DataLoader(dev, batch_size=100, shuffle=True)
te_loader = DataLoader(te, batch_size=100, shuffle=True)

model = BertVec("../../bert_en/")
s_model = LSTMVec("../saved/bert_embedding.pkl", "../../bert_en/", num_layers=2)
trainer = Distill_LM_Trainer(model, s_model, model.bert.config)
trainer.DistillTrain(tr_loader, learning_rate=2e-5, max_epochs=10, model_file="Bert_Distill_LSTM.pkl")
